#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 导入所需的库
from keras.datasets import mnist  # 从 Keras 库中导入 MNIST 数据集
import matplotlib.pyplot as plt  # 导入 Matplotlib 的 pyplot 模块，用于数据可视化
from keras import models  # 从 Keras 库中导入 models 模块，用于构建神经网络模型
from keras import layers  # 从 Keras 库中导入 layers 模块，用于向模型添加不同类型的层
import numpy as np  # 导入 Numpy 库，用于数组操作和数值计算
from tensorflow.keras.utils import to_categorical  # 从 TensorFlow 的 Keras 工具模块中导入 to_categorical 函数，用于将类别标签转换为 one-hot 编码


# 定义模型
def model_conv():
    # 创建一个顺序模型，按顺序添加层
    model = models.Sequential()
    # 添加一个 2D 卷积层，32 是滤波器的数量，(3, 3) 是滤波器的大小，activation='relu' 使用 ReLU 激活函数，
    # input_shape=(28, 28, 1) 是输入图像的形状（28x28 像素，1 个通道，因为是灰度图像）
    model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
    # 添加一个 2x2 的最大池化层，用于减少特征图的尺寸
    model.add(layers.MaxPooling2D((2, 2)))
    # 添加另一个 2D 卷积层，64 个滤波器，使用 ReLU 激活函数
    model.add(layers.Conv2D(64, (3, 3), activation='relu'))
    # 添加另一个 2x2 的最大池化层
    model.add(layers.MaxPooling2D((2, 2)))
    # 添加第三个 2D 卷积层，64 个滤波器，使用 ReLU 激活函数
    model.add(layers.Conv2D(64, (3, 3), activation='relu'))
    # 将特征图展平为一维向量，以便连接到全连接层
    model.add(layers.Flatten())
    # 添加一个具有 64 个神经元的全连接层，使用 ReLU 激活函数
    model.add(layers.Dense(64, activation='relu'))
    # 添加一个具有 10 个神经元的全连接层，使用 softmax 激活函数，用于多分类任务（MNIST 有 10 个类别，0-9 的数字）
    model.add(layers.Dense(10, activation='softmax'))
    # 编译模型，使用 RMSprop 优化器，分类交叉熵作为损失函数，评估指标为准确率
    model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['acc'])
    return model


# 导入 MNIST 数据集
# 从 Keras 中加载 MNIST 数据集，得到训练集和测试集的数据和标签
(train_data, train_labels), (test_data, test_labels) = mnist.load_data()
# 打印训练集的数据和标签的形状
print('train_shape {} {}'.format(train_data.shape, train_labels.shape))
# 打印测试集的数据和标签的形状
print('test_shape {} {}'.format(test_data.shape, test_labels.shape))


# 数据预处理
# 将训练集数据重塑为 4D 张量，符合卷积层的输入形状，添加通道维度（灰度图像，通道数为 1）
x_train = train_data.reshape((60000, 28, 28, 1))
# 将训练集数据转换为 float32 类型并归一化到 0 到 1 的范围
x_train = x_train.astype('float32') / 255
# 将测试集数据重塑为 4D 张量
x_test = test_data.reshape((10000, 28, 28, 1))
# 将测试集数据转换为 float32 类型并归一化
x_test = x_test.astype('float32') / 255
# 将训练集的标签转换为 one-hot 编码，方便模型处理多分类问题
y_train = to_categorical(train_labels)
# 将测试集的标签转换为 one-hot 编码
y_test = to_categorical(test_labels)
# 打印处理后训练集数据和标签的形状
print(x_train.shape, y_train.shape)


# 定义模型
# 调用 model_conv 函数创建模型
model = model_conv()
# 打印模型的结构信息，包括层的信息、参数数量等
print(model.summary())


# 开始训练
# 使用训练集数据和标签训练模型，训练 5 个轮次，批大小为 64，使用 10% 的训练数据作为验证集
his = model.fit(x_train, y_train, epochs=5, batch_size=64, validation_split=0.1)


# 计算准确度
# 使用测试集数据和标签评估模型，计算损失和准确率
loss, acc = model.evaluate(x_test, y_test)
# 打印测试集上的损失和准确率
print('loss {}, acc {}'.format(loss, acc))


# 保存模型
# 将训练好的模型保存为 HDF5 文件，方便后续加载和使用
model.save("my_mnist_model.h5")